//////////////////////////////////////////////////////////////////////////
//
// This file is a part of UKROP project.
// Copyright (c) Vladimir Gumenuk
//
// PROJECT:         Xort Custom Debug Engine
// FILE:            StackFrame.cpp : Implementation of CStackFrame
// AUTHOR:          Vladimir Gumenuk
// DATE:            Dec-2012 -- Jan-2013
// MODIFICATIONS:
//
//////////////////////////////////////////////////////////////////////////
 
#include "stdafx.h"
#include "xort.h"
#include "xruntime.h"
#include "DebugProperty.h"
#include "DebugStack.h"
#include "StackFrame.h"
#include "Expression.h"
#include "dllmain.h"

//////////////////////////////////////////////////////////////////////////
// CStackFrame
CStackFrame::CStackFrame()
{

}

CStackFrame::~CStackFrame()
{
    ATLTRACE( "CStackFrame::~CStackFrame()\n" );
}

HRESULT CStackFrame::FinalConstruct()
{
    return S_OK;
}

void CStackFrame::FinalRelease()
{
    UnInit();
}

STDMETHODIMP CStackFrame::Init( IDebugDocumentContext2 *pDocContext, IDebugProgram2 *pProgram, UINT64 addrMin, UINT64 addrMax, FRAMEINFO *pFrameInfo )
{
    m_spDocCtx   = pDocContext;
    m_spProgram  = pProgram;
    m_addrMin    = addrMin;
    m_addrMax    = addrMax;
    
    m_pFrameInfo = (CFrameInfo *)pFrameInfo;
    m_pFrameInfo->AddRef();
    return S_OK;
}

STDMETHODIMP CStackFrame::PushPropertyList()
{
    LPPROPLIST pList = new PROPERTIES_LIST();
    if( !pList )
        return E_OUTOFMEMORY;

    if( m_stackProps.count() > 0 )
    {
        LPPROPLIST pTopList = m_stackProps.top();
        pTopList->copy( *pList ); // copy old frame locals to new frame, to update them later
    }

    m_stackProps.push( pList );
    return S_OK;
}

STDMETHODIMP CStackFrame::PopPropertyList()
{
    if( m_stackProps.count() > 0 )
        m_stackProps.pop();
    return S_OK;
}

STDMETHODIMP CStackFrame::AddProperty( IDebugProperty2 *pProperty )
{
    LPPROPLIST pTopList = m_stackProps.top();

    CComQIPtr<IXProperty> spXProp( pProperty );
    Symbol *pSym = NULL;

    if( SUCCEEDED( spXProp->GetSymbol( (IUnknown**)&pSym ) ) )
    {
        PROPERTIES_LIST::iterator theIter = pTopList->begin();
        IDebugProperty2 *pProp = theIter == NULL ? NULL : (IDebugProperty2*)theIter;

        while( pProp )
        {
            Symbol *pLocal = NULL;

            CComQIPtr<IXProperty> spXLocal( pProp );
            if( SUCCEEDED( spXLocal->GetSymbol( (IUnknown**)&pLocal ) ) )
            {
                if( pSym->m_iIndex == pLocal->m_iIndex ) // means that this symbol is overridden by new symbol in new frame
                {
                    pTopList->erase( theIter );
                    break;
                }
            }

            theIter++;
            pProp = theIter == NULL ? NULL : (IDebugProperty2*)theIter;
        }
    }

    pTopList->push_back( pProperty );
    return S_OK;
}

STDMETHODIMP CStackFrame::UnInit() // this call removes internal X-references
{
    if( m_pFrameInfo )
    {
        CFrameInfo::Clear( *m_pFrameInfo );
        m_pFrameInfo->Release();
        m_pFrameInfo = NULL;
    }

    m_stackProps.clear();
    return S_OK;
}

////////////////////////////////////////////////////////////////////////////// 
// IDebugStackFrame2 
STDMETHODIMP CStackFrame::GetCodeContext( IDebugCodeContext2 **ppCodeCxt ) 
{
    //ATLTRACE( "CStackFrame::GetCodeContext()\n" );
    return m_spDocCtx->QueryInterface( IID_IDebugCodeContext2, (void**)ppCodeCxt ); 
}

STDMETHODIMP CStackFrame::GetDocumentContext(IDebugDocumentContext2 **ppCxt) 
{ 
    //ATLTRACE( "CStackFrame::GetDocumentContext()\n" );
    return m_spDocCtx->QueryInterface( IID_IDebugDocumentContext2, (void**)ppCxt ); 
}

STDMETHODIMP CStackFrame::GetName( BSTR* pbstrName ) 
{
    if( !pbstrName )
        return E_INVALIDARG;
    *pbstrName = SysAllocString( m_pFrameInfo->m_bstrFuncName );
    //ATLTRACE( "CStackFrame::GetName() -> %S\n", m_pFrameInfo->m_bstrFuncName );
    return S_OK;
}

STDMETHODIMP CStackFrame::GetInfo( FRAMEINFO_FLAGS dwFieldSpec, UINT nRadix, FRAMEINFO *pFrameInfo ) 
{
    if( !pFrameInfo )
        return E_INVALIDARG;
    // ignore request, just return all that we have :)
    CFrameInfo::Copy( *pFrameInfo, *m_pFrameInfo );
    //ATLTRACE( "CStackFrame::GetInfo()\n" );
    return S_OK;
}

STDMETHODIMP CStackFrame::GetPhysicalStackRange( UINT64* paddrMin, UINT64* paddrMax )
{
    //ATLTRACE( "CStackFrame::GetPhysicalStackRange()\n" );
    if( paddrMin )
        *paddrMin = m_addrMin;
    if( paddrMax )
        *paddrMax = m_addrMax;
    return S_OK;
}

STDMETHODIMP CStackFrame::GetExpressionContext( IDebugExpressionContext2 **ppExprCxt )
{
    if( !ppExprCxt )
        return E_INVALIDARG;
    return this->QueryInterface( IID_IDebugExpressionContext2, (void**)ppExprCxt );
}

STDMETHODIMP CStackFrame::GetLanguageInfo( BSTR* pbstrLanguage, GUID* pguidLanguage ) 
{
    //ATLTRACE( "CStackFrame::GetLanguageInfo()\n" );
    return m_spDocCtx->GetLanguageInfo( pbstrLanguage, pguidLanguage );
}

HRESULT CStackFrame::GetDebugProperty( IDebugProperty2 **ppDebugProp ) 
{
    ATLTRACE( "CStackFrame::GetDebugProperty()\n" );
    return E_NOTIMPL;
}

STDMETHODIMP CStackFrame::EnumProperties( DEBUGPROP_INFO_FLAGS dwFields, UINT nRadix, REFGUID guidFilter, 
    DWORD dwTimeout, ULONG *pcelt, IEnumDebugPropertyInfo2 **ppepi ) 
{
    //ATLTRACE( "CStackFrame::EnumProperties()\n" );

    CComPtr<IXEnumPropertyInfo> spXEnum;
    HRESULT hr = CEnumDebugPropertyInfo::Create( &spXEnum );
    if( FAILED( hr ) )
        return hr;

    //if( guidFilter == guidFilterRegisters ) // XVM doesn't have registers, I will just return raw stack in this case
    //{
    //    hr = CreateRegistersInfo( spXEnum, nRadix );
    //} else { // locals

        LPPROPLIST pTopList = m_stackProps.top();
        PROPERTIES_LIST::iterator theIter = pTopList->begin();
        IDebugProperty2* pProp = ( theIter == NULL ) ? NULL : (IDebugProperty2*)theIter;

        while( pProp )
        {
            CPropertyInfo *pDbgProp = new CPropertyInfo();
            pDbgProp->AddRef();

            if( SUCCEEDED( hr = pProp->GetPropertyInfo( dwFields, nRadix, dwTimeout, NULL, 0, pDbgProp ) ) )
            {
                spXEnum->AddProperty( pDbgProp );
            }

            pDbgProp->Release();

            theIter++;
            pProp = ( theIter == NULL ) ? NULL : (IDebugProperty2*)theIter;
        }

        *pcelt = pTopList->size();
    //}

    if( SUCCEEDED( hr ) )
        hr = spXEnum->QueryInterface( IID_IEnumDebugPropertyInfo2, (void**)ppepi );

    return hr;
}

STDMETHODIMP CStackFrame::GetThread( IDebugThread2** ppThread ) 
{
    //ATLTRACE( "CStackFrame::GetThread()\n" );
    return m_spProgram->QueryInterface( IID_IDebugThread2, (void**)ppThread );
} 

STDMETHODIMP CStackFrame::ParseText( LPCOLESTR pszCode, PARSEFLAGS dwFlags, UINT nRadix, IDebugExpression2 **ppExpr, BSTR *pbstrError, UINT *pichError )
{
    if( !pszCode || !ppExpr )
        return E_INVALIDARG;

    CComObject<CExpression> *p;
    HRESULT hr = CComObject<CExpression>::CreateInstance( &p );
    if( SUCCEEDED( hr ) )
    {
        p->AddRef();

        if( SUCCEEDED( hr = p->Initialize( pszCode, (IUnknown*)&m_stackProps, (DWORD)m_addrMin, nRadix ) ) )
        {
            hr = p->QueryInterface( IID_IDebugExpression2, (void**)ppExpr );
        }

        p->Release();
    }

    return hr;
}

HRESULT CStackFrame::CreateRegistersInfo( IXEnumPropertyInfo *pXEnum, UINT nRadix )
{
    HRESULT hr = S_OK;

    CComPtr<IXProgram> spXProg = _AtlModule.getXProgramNoAddRef();
    XRuntime *pRT = NULL;
    spXProg->GetXRuntime( (IUnknown**)&pRT );
    XCELLSARRAY& stack = pRT->get_Stack();
    dword dwStackBase  = pRT->get_StackBase();
    
    XortElement& sysFrame = stack.get( dwStackBase );
    dword dwArgsCount = sysFrame.bType;

    LPPROPLIST pTopList = m_stackProps.top();
    dword dwNumLocals = pTopList->size();
    dword dwTotal = dwNumLocals + dwArgsCount + 1; // + 1 system frame
    int iIdx = dwArgsCount ? -( (int)dwArgsCount ) : 0 ;
    dword dwStartIdx = dwStackBase - dwArgsCount;
    
    CComPtr<IXProperty> spXPopertyRoot;
    if( SUCCEEDED( CDebugProperty::Create( &spXPopertyRoot ) ) )
    {
        spXPopertyRoot->Initialize( NULL, 0, NULL, NULL, DBG_PROP_STKREG_ROOT );

        while( dwTotal-- )
        {
            if( stack.isValidIndex( dwStartIdx ) )
            {
                XortElement& local = stack.get( dwStartIdx );

                CComPtr<IXProperty> spXPopertyReg;
                if( SUCCEEDED( CDebugProperty::Create( &spXPopertyReg ) ) )
                {
                    CPropertyInfo *pPorp = new CPropertyInfo();
                    pPorp->AddRef();
                    spXPopertyReg->Initialize( NULL, dwStartIdx, NULL, NULL, DBG_PROP_STKREG );

                    CComQIPtr<IDebugProperty2> spDbgProp( spXPopertyReg );
                    spXPopertyRoot->AddProperty( spDbgProp );

                    pPorp->Release();
                }
            }

            dwStartIdx++;
            iIdx++;
        }

        CPropertyInfo *pPorpRoot = new CPropertyInfo();
        pPorpRoot->AddRef();
        CComQIPtr<IDebugProperty2> spDbgPropRoot( spXPopertyRoot );
        spDbgPropRoot->GetPropertyInfo( DEBUGPROP_INFO_STANDARD, nRadix, 0, NULL, 0, pPorpRoot );
        pXEnum->AddProperty( pPorpRoot );
        pPorpRoot->Release();
    }

    return hr;
}

//////////////////////////////////////////////////////////////////////////
//
CEnumFrameInfo::CEnumFrameInfo()
{

}

CEnumFrameInfo::~CEnumFrameInfo()
{

}

HRESULT CEnumFrameInfo::FinalConstruct()
{
    return S_OK;
}

void CEnumFrameInfo::FinalRelease()
{
    m_listFrames.clear();
}

STDMETHODIMP CEnumFrameInfo::AddFrameInfo( FRAMEINFO *pFrameInfo )
{
    if( !pFrameInfo )
        return E_INVALIDARG;

    CFrameInfo *pFrame = (CFrameInfo *)pFrameInfo;
    m_listFrames.push_front( pFrame ); // will ref+ object
    m_theIter = m_listFrames.begin();
    //ATLTRACE( "CEnumFrameInfo::AddFrameInfo() -> size: %d\n", m_listFrames.size() );
    return S_OK;
}

////////////////////////////////////////////////////////////////////////////// 
// IEnumDebugFrameInfo2 
STDMETHODIMP CEnumFrameInfo::Skip( ULONG celt ) 
{
    //ATLTRACE( "CEnumFrameInfo::Skip() -> %d\n", celt );
    m_theIter += celt;
    return S_OK;
}

STDMETHODIMP CEnumFrameInfo::Clone( IEnumDebugFrameInfo2** ppEnum ) 
{
    ATLTRACE( "CEnumFrameInfo::Clone()\n" );
    return E_NOTIMPL;
}

STDMETHODIMP CEnumFrameInfo::Reset() 
{
    m_theIter = m_listFrames.begin();
    return S_OK;
} 

STDMETHODIMP CEnumFrameInfo::GetCount(ULONG* pcelt) 
{
    if( !pcelt )
        return E_INVALIDARG;
    *pcelt = m_listFrames.size();
    return S_OK; 
} 

STDMETHODIMP CEnumFrameInfo::Next( ULONG celt, FRAMEINFO *rgelt, ULONG* pceltFetched ) 
{ 
    //ATLTRACE( "CEnumFrameInfo::Next() -> %d\n", celt );
    if( !rgelt || celt == 0 ) //|| ( celt != 1 && pceltFetched == NULL ) )
        return E_INVALIDARG;

    if( m_theIter == NULL )
    {
        if( pceltFetched )
            *pceltFetched = 0;
        return S_FALSE;
    }
    
    ULONG lRequested  = celt;
    ULONG lNumFetched = 0;

    while( celt-- )
    {
        CFrameInfo::Copy( *rgelt, *((CFrameInfo*)m_theIter) );
        rgelt++;
        m_theIter++;
        lNumFetched++;
        if( m_theIter == NULL )
            break;
    }

    if( pceltFetched )
        *pceltFetched = lNumFetched;

    //ATLTRACE( "CEnumFrameInfo::Next() -> requested: %d, fetched: %d\n", lRequested, lNumFetched );
    return lRequested == lNumFetched ? S_OK : S_FALSE;
}

//////////////////////////////////////////////////////////////////////////
// StackFrame.cpp : Implementation of CEnumDebugPropertyInfo
CEnumDebugPropertyInfo::CEnumDebugPropertyInfo()
{

}

CEnumDebugPropertyInfo::~CEnumDebugPropertyInfo()
{

}

HRESULT CEnumDebugPropertyInfo::FinalConstruct()
{
    return S_OK;
}

void CEnumDebugPropertyInfo::FinalRelease()
{
}

STDMETHODIMP CEnumDebugPropertyInfo::AddProperty( DEBUG_PROPERTY_INFO *pProp ) // actually internal CPropertyInfo passed!
{
    if( !pProp )
        return E_INVALIDARG;
    CPropertyInfo *p = (CPropertyInfo *)pProp;
    m_listProperties.push_back( p ); // +ref
    m_theIter = m_listProperties.begin();
    return S_OK;
}

STDMETHODIMP CEnumDebugPropertyInfo::Next( ULONG celt, DEBUG_PROPERTY_INFO *rgelt, ULONG *pceltFetched )
{
    if( !rgelt || celt == 0 ) //|| ( celt != 1 && pceltFetched == NULL ) )
        return E_INVALIDARG;

    if( m_theIter == NULL )
    {
        if( pceltFetched )
            *pceltFetched = 0;
        return S_FALSE;
    }

    ULONG lRequested  = celt;
    ULONG lNumFetched = 0;

    while( celt-- )
    {
        CPropertyInfo *pProp = (CPropertyInfo *)m_theIter;
        pProp->Copy( rgelt );
        rgelt++;
        m_theIter++;
        lNumFetched++;
        if( m_theIter == NULL )
            break;
    }

    if( pceltFetched )
        *pceltFetched = lNumFetched;

    return lRequested == lNumFetched ? S_OK : S_FALSE;
}

STDMETHODIMP CEnumDebugPropertyInfo::Skip( ULONG celt )
{
    m_theIter += celt;
    return S_OK;
}

STDMETHODIMP CEnumDebugPropertyInfo::Reset()
{
    m_theIter = m_listProperties.begin();
    return S_OK;
}

STDMETHODIMP CEnumDebugPropertyInfo::Clone( IEnumDebugPropertyInfo2 **ppEnum )
{
    return E_NOTIMPL;
}

STDMETHODIMP CEnumDebugPropertyInfo::GetCount( ULONG *pcelt )
{
    if( !pcelt )
        return E_INVALIDARG;
    *pcelt = m_listProperties.size();
    return S_OK;
}
